#!/usr/bin/env python
# encoding: UTF-8
"""
A script to merge water mediated contacts in the database generated
by contacts.py. If the contact to water has no counterpart, it is removed.

After that, a removal of duplicate contacts is performed.
"""
__version__ = "$Id$"

import sys
import sqlite3
import logging

log = logging.getLogger("db_cleanup")
log.addHandler(logging.StreamHandler())
log.setLevel(logging.INFO)

def get_tables(database_name):
    """
    Generates of pairs of tables in <database_name> database.
    Returns list of tuples in the form:
    [ (u'contacts_9ZZZ', u'atoms_9ZZZ')
       .
       .
       . ]
       
    
    Beware - if the database is corrupted, there is no guarantee that
    the atoms table exists.
    """
    conn = sqlite3.connect(database_name)
    cur = conn.cursor()
    ret = []
    try:
        tables = cur.execute("""SELECT name FROM sqlite_master
                    WHERE type='table' AND name LIKE 'contacts_%'
                    ORDER BY name""")
        for row in tables:
            pdb = row[0].split("_")[1]
            ret.append((row[0], "atoms_%s" % pdb))
    finally:
        conn.close()
        
    return ret

def merge(contacts_table, atoms_table, db_name):
    conn = sqlite3.connect(db_name)
    cur1 = conn.cursor()
    
    cur1.executescript("""
    PRAGMA temp_store = MEMORY;
    PRAGMA cache_size = 80000;
    PRAGMA synchronous=0;
    CREATE INDEX IF NOT EXISTS "acceptor_id_type_%(contacts)s" ON "%(contacts)s" ("acceptor_id", "type");
    CREATE INDEX IF NOT EXISTS "donor_id_type_%(contacts)s" ON "%(contacts)s" ("donor_id", "type");
    CREATE INDEX IF NOT EXISTS "donor_id_acceptor_%(contacts)s" ON "%(contacts)s" ("donor_id", "acceptor_id");
    """ % {'contacts' : contacts_table, 'atoms' : atoms_table})
    conn.commit()
    
    cur2 = conn.cursor()
    cur3 = conn.cursor()
    log.info("          merging wm contacts: %s" % contacts_table)
    # najde všechny kontakty, co mají jako donor vodu

    merge_SQL = """
INSERT INTO %(contacts)s
    (donor_id, acceptor_id, directness, type, cat_a, cat_d)
SELECT donor, acceptor, directness, type, cat_a, cat_d FROM (
    select
        c1.donor_id AS donor, c2.acceptor_id AS acceptor, 'W' as directness, c1.type AS type, c1.cat_d AS cat_d, c2.cat_a AS cat_a
    from
        %(contacts)s c1, %(contacts)s c2, %(atoms)s a
    where
        c1.acceptor_id = c2.donor_id
        and c1.acceptor_id=a.id
        and a.amino='HOH'
        and c1.type = c2.type
    UNION ALL
    select
        c1.donor_id AS donor, c2.donor_id AS acceptor, 'X' as directness, c1.type AS type, c1.cat_d AS cat_d, c2.cat_d AS cat_a
    from
        %(contacts)s c1, %(contacts)s c2, %(atoms)s a
    where
        c1.acceptor_id = c2.acceptor_id
        and c1.acceptor_id=a.id
        and c1.type = c2.type
        and a.amino='HOH'
        and c1.donor_id < c2.donor_id
    UNION ALL
    select
        c1.acceptor_id AS donor, c2.acceptor_id AS acceptor, 'X' as directness, c1.type AS type, c1.cat_a AS cat_d, c2.cat_a AS cat_a
    from
        %(contacts)s c1, %(contacts)s c2, %(atoms)s a
    where
        c1.donor_id = c2.donor_id
        and c1.donor_id=a.id
        and c1.type = c2.type
        and a.amino='HOH'
        and c1.acceptor_id < c2.acceptor_id
    ), %(atoms)s d, %(atoms)s a
    WHERE
            d.id = donor        -- join constraint
        AND a.id = acceptor     -- join constraint
        AND (
          (d.amino IN ("A", "G", "C", "T", "U")          -- donor is DNA
            AND a.amino NOT IN ("A", "G", "C", "T", "U") -- acceptor is protein
          ) 
          OR
          (a.amino IN ("A", "G", "C", "T", "U")          -- acceptor is DNA
            AND d.amino NOT IN ("A", "G", "C", "T", "U") -- donor is protein
          )
        )
    """

    cur2.execute(merge_SQL % {'contacts' : contacts_table, 'atoms' : atoms_table})
                           
    conn.commit()
    
    delete_waters(cur1, contacts_table, atoms_table)
    conn.commit()    

def delete_waters(cursor, contacts_table, atoms_table):
    """Deletes all contacts to water in the db"""
    cursor.execute("""DELETE FROM %(contacts)s
                        WHERE donor_id IN
                            (SELECT id FROM %(atoms)s WHERE amino = "HOH")
                          OR acceptor_id IN
                            (SELECT id FROM %(atoms)s WHERE amino = "HOH")
                        """ % {'contacts' : contacts_table,
                               'atoms'    : atoms_table})
    cursor.execute("""DELETE FROM %s WHERE amino = "HOH"
                   """ % atoms_table)
                   
def remove_duplicate_contacts(contacts_table, db_name):
    """Removes duplicate (having the same donor and acceptor) contacts.
    It tries to preserve hydrogen bonds over near contacts. The second
    criterion is the directness of the bond, i.e. the direct contact is
    preserved, while the water-mediated contact is discarded.
    As a result, if there is a water-mediated hydrogen contact
    and a direct near contact, the hydrogen contact will survive.
    """
    
    conn = sqlite3.connect(db_name)
    conn.executescript("""PRAGMA temp_store = MEMORY;
                           PRAGMA cache_size = 80000;
                           PRAGMA synchronous=0;""")
    cur1 = conn.cursor()
    cur2 = conn.cursor()
    cur3 = conn.cursor()
    log.info("          removing duplicates: %s" % contacts_table)
    
    # updates near contacts for which there is a same contact
    # (hydrogen or near), which has only a donor and acceptor
    # switched (near contacts don't distinguish between acceptor and
    # donor in the sense the hydrogen bond does), so that duplicates
    # can be properly removed in next step
    cur1.execute("SELECT donor_id, acceptor_id FROM %s" % contacts_table)
    for x in cur1:
        cur2.execute("""UPDATE %s SET donor_id = ?, acceptor_id = ?
                    WHERE donor_id = ? AND acceptor_id = ? AND type = ?""" % contacts_table,
                    (x[1], x[0], x[0], x[1], "N"))

    cur1.execute("""SELECT donor_id, acceptor_id, COUNT(*) 
                    FROM %s GROUP BY donor_id, acceptor_id 
                    HAVING count(*) > 1""" % contacts_table)
                    
    for row in cur1:
        donor_id = row[0]
        acceptor_id = row[1]
        
        log.debug("%s, %s" % (donor_id, acceptor_id))

        # ordered by directness, type
        # (Hydrogen < Near, Direct < Water-mediated)
        cur2.execute("""SELECT id, type, directness FROM %s
                        WHERE donor_id = ? AND acceptor_id = ?
                        ORDER BY directness, type""" % contacts_table,
                        (donor_id, acceptor_id))
        to_delete = cur2.fetchall()[1:]
        
        assert len(to_delete) > 0
        
        for c in to_delete:
            cur3.execute("DELETE FROM %s WHERE id = ?" % contacts_table, (c[0], ))
    
    conn.commit()
    conn.close()

def main(argv):
    if len(argv) < 2:
        print "Usage: %s db_name" % sys.argv[0]
        return 1
    
    if "-d" in argv:
        log.setLevel(logging.DEBUG)
        argv.remove("-d")
        
    DB_NAME = argv[1]
    
    tables = get_tables(DB_NAME)
    num_tables = len(tables)
    
    for i, (contacts, atoms) in enumerate(tables):
        merge(contacts, atoms, DB_NAME)
        n = remove_duplicate_contacts(contacts, DB_NAME)
        log.info("(%3d/%3d) finished processing of %s" % (i, num_tables, contacts))
        
    conn = sqlite3.connect(DB_NAME)
    log.info("sqlite> VACUUM;")
    conn.execute("VACUUM")
    conn.commit()
    conn.close()
    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))
